
import time

from HDetectionv3 import hdetect
from LDetection2 import ldetect2
import json
import cv2
# Social Distance relate
import AI_Algorithims.SocialDisAnalyser.analyser as analyser
from yolo.tool.torch_utils import do_detect
from Websocket.tornado_websocket_server import LandArg, SocialDisArg
from ws4py.client.threadedclient import WebSocketClient
from HDetection import create_model
from Websocket.tornado_websocket_server import is_json
import json

from yolo.tool.utils import load_class_names


SDArgs = SocialDisArg(angle_factor=0.8, H_zoom_factor=1.2, confid=0.3)
arg_single = LandArg(True, 544, 3, 0.1, 0.3, "")

class Client(WebSocketClient):
    def opened(self):

        self.model = create_model()
        self.classnames = load_class_names("D:/PYProject/FlaskServer/yolo/data/mark.names")

        self.curr_height = 1000
        self.rtmp_thread = ""
        self.rtmp_server = 'rtmp://8.134.15.236:1935/live/'
        self.mission_type = "land"

        self.net = analyser.getSocialDisNet()
        self.labels = analyser.getLabels()
        self.device = analyser.getDevice()
        self.join_msg = {"cmd": "join", "roomId": "MEC", "uid": "MecClient"}
        self.send(json.dumps(self.join_msg))


    def closed(self, code, reason=None):
        print("Closed down", code, reason)


    # 服务器返回消息
    def received_message(self, message):
        print(len(str(message)))
        value = message
        print(value)
        if is_json(str(value)):
            j_value = json.loads(str(value))
            try:
                if j_value["cmd"] == "mecdata":
                    self.cap = cv2.VideoCapture(self.rtmp_server + j_value["recordName"])
                    # 调用cv2方法获取cap的视频帧（帧：每秒多少张图片）
                    self.fps = self.cap.get(cv2.CAP_PROP_FPS)
                    print(self.fps)
                    # 获取cap视频流的每帧大小
                    self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                    self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                    self.half_width = self.width / 2
                    self.half_height = self.height / 2
                    self.size = (self.width, self.height)
                    print(self.size)

                    ret, image = self.cap.read()
                    count = 0
                    (H, W) = image.shape[:2]

                    while ret:
                        if not ret:
                            break
                        if count % 6 == 0 or count == 0:
                            print(count)
                            ###Social DIS Block###------START
                            image, rect_list, line_list = analyser.socialDisDetecter(self.net, self.labels, image,
                                                                                     self.device, H, W, SDArgs)
                            rect_list = [i for item in rect_list for i in item]
                            line_list = [i for item in line_list for i in item]
                            if len(rect_list) == 0:
                                rect_list = [0]
                            if len(line_list) == 0:
                                line_list = [0]
                            msg = {"cmd": "mecdata", "remoteRoomId": j_value["roomId"], "remoteUid": j_value["uid"],
                                   "uid": "MecClient",
                                   "type": 1, "msg": "SD", "time": 1, "deviation": count,
                                   "rect_list": str(rect_list), "line_list": str(line_list)}
                            print(msg)
                            self.send(json.dumps(msg))
                            ###Social DIS Block###------END

                            ###CV Landing Block###------START
                            sized = cv2.resize(image, (arg_single.size, arg_single.size))
                            sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
                            start = time.clock()
                            boxes = do_detect(self.model, sized, arg_single.conf, arg_single.nms,
                                              arg_single.use_cuda)
                            end = time.clock()
                            print((start - end), "---------------------------------------------")
                            if len(boxes[0]) == 4 or len(boxes[0]) == 5 or len(boxes[0]) == 6:
                                cmd, image, dx, dy, angle = ldetect2(boxes, image, self.half_width,
                                                                     self.half_height)
                            else:
                                angle = 0
                                cmd, image, dx, dy = hdetect(boxes, image=image, mclass_names=self.classnames)
                            if cmd == "null":
                                msg = {"cmd": "mecdata", "remoteRoomId": j_value["roomId"], "remoteUid": j_value["uid"],
                                       "uid": "MecClient", "type": 2, "msg": cmd, "time": 1, "deviation": count,
                                       "rect_list": str(rect_list), "line_list": str(line_list)}
                                print(msg)
                                self.send(json.dumps(msg))
                            else:
                                value = str({"dx": dx, "dy": dy, "mpitch": round(dy / self.half_height, 2),
                                             "mroll": round(dx / self.half_width, 2), "myaw": angle,
                                             "mthrottle": 0})
                                print(value)
                                msg = {"cmd": "mecdata", "remoteRoomId": j_value["roomId"], "remoteUid": j_value["uid"],
                                       "uid": "MecClient", "type": 2, "msg": value, "time": 1, "deviation": count,
                                       "rect_list": str(rect_list), "line_list": str(line_list)}
                                print(msg)
                                self.send(json.dumps(msg))
                            print("before", dx, dy)
                            ###CV Landing Block###------END

                            cv2.namedWindow("Landing", 0)
                            cv2.imshow("Landing", image)
                            cv2.waitKey(1)
                        count = count + 1
                        ret, image = self.cap.read()

            except ValueError as e:
                print(e)

if __name__ == '__main__':
    try:
        # 服务器连接地址wss://real.okex.com:10440/websocket/okexapi  (该地址需翻墙，测试可找其他地址)
        ws = Client('ws://8.134.15.236:3000')
        ws.connect()
        ws.run_forever()
    except KeyboardInterrupt:
        ws.close()